# encoding = utf-8
import json
import sys

import aioredis
from loguru import logger
import uvicorn
from fastapi import FastAPI, Query, WebSocket, WebSocketDisconnect, HTTPException, Request, Form, Body
from fastapi.middleware.cors import CORSMiddleware
import jwt
from datetime import datetime, timedelta
from functools import wraps
from fastapi.responses import JSONResponse
from fastapi.background import BackgroundTasks
from config import DOMAIN, SECRET_KEY, TriggerExecuteToken
from playwright_method import safe_globals, send_message_to_redis
from api_method import verify_permissions_, login_, register_, check_email_, send_email, email_reset_password_, \
    reset_password_, get_all_method_, add_project_, get_project_, update_project_, delete_project_, add_uicase_, \
    update_uicase_, update_uicase_state_, delete_uicase_, delete_uicase_step_, get_project_uicase_, copy_uicase_, \
    add_apicase_, update_apicase_, delete_apicase_, delete_apicase_step_, get_project_apicase_, copy_apicase_, \
    add_test_report_, select_case_id, background_run_case, get_test_report_, delete_test_report_, set_scheduled_tasks_, \
    get_scheduled_tasks_, add_account_, update_account_, get_account_, delete_account_, copy_account_, \
    trigger_execute_case_, verify_sign_

app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有来源
    allow_credentials=True,
    allow_methods=["*"],  # 允许所有HTTP方法，包括OPTIONS
    allow_headers=["*"],
)

domain = DOMAIN
# 配置日志
logger.remove()
# 添加文件处理器，设置日志文件路径
logger.add("log/app.log", level="INFO", rotation="10 MB", retention="30 days", compression="zip")
logger.add(sys.stderr, level="INFO")  # 输出到控制台


# 自定义错误响应格式
def error_response(status_code: int, message: str, request: Request) -> dict:
    return {
        "code": status_code,
        "mes": message,
        "path": request.url.path,
        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    }


# HTTP异常处理器
@app.exception_handler(HTTPException)
async def http_exception_handler(request: Request, exc: HTTPException):
    return JSONResponse(
        status_code=exc.status_code,
        content=error_response(exc.status_code, exc.detail, request)
    )


# 全局异常处理器
@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    return JSONResponse(
        status_code=500,
        content=error_response(500, "服务器内部错误", request)
    )


def verify_request_sign(f):
    @wraps(f)  # 将原函数的信息复制到包装函数中 避免信息丢失
    # *args 它可以传递任意数量的参数
    # **kwargs 允许使用没有事先定义的参数名
    async def decorated_function(request: Request, *args, **kwargs):
        # 获取请求头中的签名和时间戳
        sign = request.headers.get('Sign') or '0'
        time = request.headers.get('Time') or '0'

        # 验证签名和时间戳
        if not verify_sign_(sign, time):  # 自定义的验证函数
            raise HTTPException(status_code=403, detail="请求非法")

        # 验证通过，继续执行原有的视图函数
        return await f(request, *args, **kwargs)

    return decorated_function


def verify_request_token(f):
    @wraps(f)
    async def decorated_function(request: Request, *args, **kwargs):
        content_type = request.headers.get('Content-Type')
        member_id = None

        if content_type:
            if content_type == 'application/json':
                body = await request.json()
                member_id = body.get('member_id')
            elif 'multipart/form-data' in content_type:
                form_data = await request.form()
                member_id = form_data.get('member_id')

        # 如果上述方法都未获取到 member_id，则尝试从查询参数获取
        if not member_id:
            member_id = request.query_params.get('member_id')

        # 获取 Authorization header
        Authorization = request.headers.get('Authorization')
        if not (member_id and Authorization):
            raise HTTPException(status_code=401, detail="member_id或Authorization未填")

        # 去掉 "Bearer " 前缀（如果有）
        token = None
        if Authorization.startswith("Bearer "):
            token = Authorization.split(" ")[1]
        else:
            token = Authorization
        if not token:
            raise HTTPException(status_code=401, detail="缺少 Token")

        try:
            # 解码并校验 JWT Token
            decoded = jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
            if int(decoded['id']) != int(member_id):
                raise HTTPException(status_code=401, detail="用户不匹配")
            if decoded['exp'] < datetime.now().timestamp():
                raise HTTPException(status_code=401, detail="Token 已过期")
            # 验证通过，继续执行原有的视图函数
            return await f(request, *args, **kwargs)

        except jwt.ExpiredSignatureError:
            raise HTTPException(status_code=401, detail="Token 已过期")
        except jwt.InvalidTokenError:
            raise HTTPException(status_code=401, detail="无效的 Token")
        except Exception as e:
            logger.error(f"{e}")
            # 如果捕获的是 HTTPException，则直接抛出
            if isinstance(e, HTTPException):
                raise e
            else:
                # 其他异常统一返回 401
                raise HTTPException(status_code=401, detail="Token 验证失败，请重新登录")

    return decorated_function


def verify_request_permissions(f):
    @wraps(f)
    async def decorated_function(request: Request, *args, **kwargs):
        content_type = request.headers.get('Content-Type')
        member_id = None

        if content_type:
            if content_type == 'application/json':
                body = await request.json()
                member_id = body.get('member_id')
            elif 'multipart/form-data' in content_type:
                form_data = await request.form()
                member_id = form_data.get('member_id')

        # 如果上述方法都未获取到 member_id，则尝试从查询参数获取
        if not member_id:
            member_id = request.query_params.get('member_id')
        if not await verify_permissions_(member_id):
            raise HTTPException(status_code=403, detail="抱歉，您暂无此权限，请向管理员申请")
        # 验证通过，继续执行原有的视图函数
        return await f(request, *args, **kwargs)

    return decorated_function


# 登录
@app.post('/api/login')
@verify_request_sign
async def login(request: Request, user: str = Form(...), pwd: str = Form(...)):
    data = await login_(user, pwd)
    if data:
        # 获取当前时间
        now = datetime.now()
        # 计算过期时间为7天后的午夜
        midnight = datetime.combine(now.date() + timedelta(days=7), datetime.max.time())
        # 生成 JWT Token
        token = jwt.encode({
            'id': data['id'],
            'exp': midnight
        }, SECRET_KEY)
        return {"id": data['id'], "user": data['user'], "token": token}
    else:
        raise HTTPException(status_code=403, detail="账号或密码错误")


# 注册
@app.post('/api/register')
@verify_request_sign
async def register(request: Request, user: str = Form(...), pwd: str = Form(...), email: str = Form(...),
                   invitationCode: str = Form(...)):
    if invitationCode != 'invitationCode':
        raise HTTPException(status_code=403, detail="邀请码错误")
    data = await register_(user, pwd, email)
    if data:
        return {'mes': data}
    else:
        raise HTTPException(status_code=403, detail="账号已存在")


# 发送重置密码邮件
@app.post('/api/checkemail')
@verify_request_sign
async def checke_mail(request: Request, email: str = Form(...)):
    data = await check_email_(email)
    if data:
        url = f'{domain}/resetpassword?id={data["id"]}&token={data["pwd"]}'
        send_email(url, email)
        return {'mes': '密码重置链接已发送至您的电子邮箱，请按照说明重置您的密码。'}
    else:
        raise HTTPException(status_code=403, detail="邮箱不存在")


# 通过邮件重置密码
@app.post('/api/email_reset_password')
@verify_request_sign
async def email_reset_password(request: Request, member_id: str = Form(...), token: str = Form(...),
                               new_pwd: str = Form(...)):
    if await email_reset_password_(member_id, token, new_pwd):
        return {'mes': '密码重置成功'}
    else:
        raise HTTPException(status_code=403, detail="账号不存在或token失效")


# 重置密码
@app.post('/api/reset_password')
@verify_request_sign
@verify_request_token
async def reset_password(request: Request, member_id: str = Form(...), old_pwd: str = Form(...),
                         new_pwd: str = Form(...)):
    if await reset_password_(member_id, old_pwd, new_pwd):
        return {'mes': '密码重置成功'}
    else:
        raise HTTPException(status_code=403, detail="旧密码错误")


# 获取所有方法
@app.get('/api/get_all_playwright')
@verify_request_sign
@verify_request_token
async def get_all_method(request: Request, case_type: int):
    data = await get_all_method_(case_type)
    return {"data": data, 'code': 200}


# 新增项目
@app.post('/api/add_project')
@verify_request_sign
@verify_request_token
@verify_request_permissions
async def add_project(request: Request, member_id: int = Body(...), name: str = Body(...),
                      project_type: int = Body(...),
                      modified_by: str = Body(...), remark: str = Body(None)):
    data = await add_project_(name, project_type, modified_by, remark)
    return {'data': data, 'code': 200}


# 获取项目列表
@app.get('/api/get_project')
@verify_request_sign
@verify_request_token
async def get_project(request: Request, member_id: str, project_id: str = None, name: str = None, page: int = None):
    data = await get_project_(name, project_id, page)
    return {"data": data, 'code': 200}


# 修改项目
@app.post('/api/update_project')
@verify_request_sign
@verify_request_token
@verify_request_permissions
async def update_project(request: Request, project_id: int = Body(...), name: str = Body(...),
                         project_type: int = Body(...), modified_by: str = Body(...), remark: str = Body(None)):
    data = await update_project_(project_id, name, project_type, modified_by, remark)
    return {'data': data, 'code': 200}


# 删除项目
@app.post('/api/delete_project')
@verify_request_sign
@verify_request_token
@verify_request_permissions
async def delete_project(request: Request, member_id: int = Body(...), project_id: int = Body(...),
                         modified_by: str = Body(...)):
    data = await delete_project_(project_id, modified_by)
    return {'data': data, 'code': 200}


# 新增uicase
@app.post('/api/add_uicase')
@verify_request_sign
@verify_request_token
@verify_request_permissions
async def add_uicase(request: Request, member_id: int = Body(...), project_id: int = Body(...),
                     project_type: int = Body(...),
                     domain: str = Body(...), title: str = Body(...), level: str = Body(...),
                     modified_by: str = Body(...),
                     steps=Body(...), state: int = Body(1), proxy: str = Body(None),
                     cookies: int = Body(0), remark: str = Body(None)):
    data = await add_uicase_(project_id, project_type, domain, title, level, modified_by, steps, state, proxy=proxy,
                             cookie=cookies,
                             remark=remark)
    return {'data': data, 'code': 200, 'mes': '新增成功'}


# 修改uicase
@app.post('/api/update_uicase')
@verify_request_sign
@verify_request_token
@verify_request_permissions
async def update_uicase(request: Request, member_id: int = Body(...), case_id: int = Body(...),
                        domain: str = Body(...), title: str = Body(...), level: str = Body(...),
                        modified_by: str = Body(...),
                        steps=Body(...), proxy: str = Body(None),
                        cookies: int = Body(0), remark: str = Body(None)):
    data = await update_uicase_(case_id, domain, title, level, modified_by, steps, proxy=proxy, cookie=cookies,
                                remark=remark)
    return {'data': data, 'code': 200, 'mes': '修改成功'}


# 修改uicase状态
@app.post('/api/update_uicase_state')
@verify_request_sign
@verify_request_token
@verify_request_permissions
async def update_uicase_state(request: Request, member_id: int = Body(...), case_id: int = Body(...),
                              state: int = Body(...),
                              modified_by: str = Body(...)):
    data = await update_uicase_state_(case_id, state, modified_by)
    return {'data': data, 'code': 200}


# 删除uicase
@app.post('/api/delete_uicase')
@verify_request_sign
@verify_request_token
@verify_request_permissions
async def delete_uicase(request: Request, member_id: int = Body(...), case_id: int = Body(...),
                        modified_by: str = Body(...)):
    data = await delete_uicase_(case_id, modified_by)
    return {'data': data, 'code': 200}


# 删除case步骤
@app.post('/api/delete_uicase_step')
@verify_request_sign
@verify_request_token
@verify_request_permissions
async def delete_uicase_step(request: Request, member_id: int = Body(...), case_id: int = Body(...),
                             step_id: int = Body(...),
                             modified_by: str = Body(...)):
    data = await delete_uicase_step_(case_id, modified_by, step_id)
    return {'data': data, 'code': 200}


# 查询项目 uicase
@app.get('/api/get_project_uicase')
@verify_request_sign
@verify_request_token
async def get_project_uicase(request: Request, member_id: int, project_id: int, case_id: int = None,
                             title: str = None, page: int = None):
    data = await get_project_uicase_(project_id, case_id, title, page)
    return {'data': data, 'code': 200}


# 复制uicase
@app.post('/api/copy_uicase')
@verify_request_sign
@verify_request_token
@verify_request_permissions
async def copy_uicase(request: Request, member_id: int = Body(...), case_id: int = Body(...),
                      project_id: int = Body(...),
                      modified_by: str = Body(...)):
    data = await copy_uicase_(case_id, project_id, modified_by)
    return {'data': data, 'code': 200}


# 添加apicase
@app.post('/api/add_apicase')
@verify_request_sign
@verify_request_token
@verify_request_permissions
async def add_apicase(request: Request, member_id: int = Body(...), project_id: int = Body(...),
                      project_type: int = Body(...),
                      domain: str = Body(...), title: str = Body(...), level: str = Body(...),
                      modified_by: str = Body(...),
                      steps=Body(...), state: int = Body(None), proxy: str = Body(None),
                      headers: int = Body(0), remark: str = Body(None)):
    data = await add_apicase_(project_id, project_type, domain, title, level, modified_by, steps, state, proxy=proxy,
                              headers=headers,
                              remark=remark)
    return {'data': data, 'code': 200}


# 修改apicase
@app.post('/api/update_apicase')
@verify_request_sign
@verify_request_token
@verify_request_permissions
async def update_apicase(request: Request, member_id: int = Body(...), case_id: int = Body(...),
                         domain: str = Body(...), title: str = Body(...), level: str = Body(...),
                         modified_by: str = Body(...),
                         steps=Body(...), proxy: str = Body(None),
                         headers: int = Body(0), remark: str = Body(None)):
    data = await update_apicase_(case_id, domain, title, level, modified_by, steps, proxy=proxy, headers=headers,
                                 remark=remark)
    return {'data': data, 'code': 200, 'mes': '修改成功'}


# 删除apicase
@app.post('/api/delete_apicase')
@verify_request_sign
@verify_request_token
@verify_request_permissions
async def delete_apicase(request: Request, member_id: int = Body(...), case_id: int = Body(...),
                         modified_by: str = Body(...)):
    data = await delete_apicase_(case_id, modified_by)
    return {'data': data, 'code': 200}


# 删除case步骤
@app.post('/api/delete_apicase_step')
@verify_request_sign
@verify_request_token
@verify_request_permissions
async def delete_apicase_step(request: Request, member_id: int = Body(...), case_id: int = Body(...),
                              step_id: int = Body(...), modified_by: str = Body(...)):
    data = await delete_apicase_step_(case_id, modified_by, step_id)
    return {'data': data, 'code': 200}


# 查询项目 apicase
@app.get('/api/get_project_apicase', )
@verify_request_sign
@verify_request_token
async def get_project_apicase(request: Request, member_id: int, project_id: int, case_id: int = None,
                              title: str = None, page: str = None):
    data = await get_project_apicase_(project_id, case_id, title, page)
    return {'data': data, 'code': 200}


# 复制apicase
@app.post('/api/copy_apicase')
@verify_request_sign
@verify_request_token
@verify_request_permissions
async def copy_apicase(request: Request, member_id: int = Body(...), case_id: int = Body(...),
                       project_id: int = Body(...),
                       modified_by: str = Body(...), ):
    data = await copy_apicase_(case_id, project_id, modified_by)
    return {'data': data, 'code': 200}


# 添加到测试报告表
@app.post('/api/add_test_report')
@verify_request_sign
@verify_request_token
async def add_test_report(request: Request, member_id: int = Body(...), case_id=Body(...),
                          title=Body(...),
                          project_name: str = Body(...), run_environment: int = Body(...)):
    data = await add_test_report_(case_id, title, project_name, run_environment)
    return {'data': data, 'code': 200, 'mes': '运行成功，请通过测试报告页面查看'}


# 运行case
@app.post('/api/run_case')
@verify_request_sign
@verify_request_token
async def run_case(request: Request, background_tasks: BackgroundTasks, member_id: int = Body(...),
                   id: int = Body(...)):
    result = await select_case_id(id)
    # 添加到后台任务
    redis = await aioredis.from_url('redis://localhost', encoding="utf-8", decode_responses=True)
    await send_message_to_redis(redis, str(id), {"status": "running", "message": "添加到后台任务"})
    background_tasks.add_task(background_run_case, result, redis, send_message_to_redis, 1)
    return {'code': 200, 'data': '已开始运行，请稍后通过测试报告页面查看结果'}


# 查询测试报告
@app.get('/api/get_test_report')
@verify_request_sign
@verify_request_token
async def get_test_report(request: Request, member_id: int, name: str = None, state: str = None, page: str = None):
    data = await get_test_report_(name, state, page)
    return {'data': data, 'code': 200}


# 删除测试报告
@app.post('/api/delete_test_report')
@verify_request_sign
@verify_request_token
@verify_request_permissions
async def delete_test_report(request: Request, member_id: int = Body(...), id=Body(...)):
    data = await delete_test_report_(id)
    return {'data': data, 'code': 200}


# 设置定时任务
@app.post('/api/set_scheduled_tasks')
@verify_request_sign
@verify_request_token
@verify_request_permissions
async def set_scheduled_tasks(request: Request, member_id: int = Body(...), project_id: int = Body(...),
                              modified_by: str = Body(...), scheduled: str = Body(None),
                              notice_type: int = Body(None), threads: int = Body(None), notice_token: str = Body(None),
                              notice_secret: str = Body(None), ):
    data = await set_scheduled_tasks_(project_id, modified_by, scheduled, notice_type, threads, notice_token,
                                      notice_secret)
    return {'data': data, 'code': 200}


# 查询定时任务
@app.get('/api/get_scheduled_tasks')
@verify_request_sign
@verify_request_token
async def get_scheduled_tasks(request: Request, member_id: int, project_id: int):
    data = await get_scheduled_tasks_(project_id)
    return {'data': data, 'code': 200}


# 新增账号
@app.post('/api/add_account')
@verify_request_sign
@verify_request_token
@verify_request_permissions
async def add_account(request: Request, member_id: int = Body(...), name: str = Body(...),
                      type: int = Body(...), cookies: str = Body(...), modified_by: str = Body(...),
                      remark: str = Body(None)):
    try:
        cookies = eval(cookies, safe_globals())
    except Exception as e:
        raise HTTPException(status_code=400, detail=f'cookie/headers格式不正确 {e}')
    data = await add_account_(name, type, json.dumps(cookies), modified_by, remark)
    return {'data': data, 'code': 200}


# 修改账号
@app.post('/api/update_account')
@verify_request_sign
@verify_request_token
@verify_request_permissions
async def update_account(request: Request, member_id: int = Body(...), id: int = Body(...), name: str = Body(...),
                         type: int = Body(...), cookies: str = Body(...), modified_by: str = Body(...),
                         remark: str = Body(None)):
    try:
        cookies = eval(cookies, safe_globals())
    except Exception as e:
        raise HTTPException(status_code=400, detail=f'cookie/headers格式不正确 {e}')
    data = await update_account_(id, name, type, json.dumps(cookies), modified_by, remark)
    return {'data': data, 'code': 200}


# 查询账号
@app.get('/api/get_account')
@verify_request_sign
@verify_request_token
async def get_account(request: Request, member_id: int, page: str = None, id: str = None, name: str = None,
                      type: int = None):
    data = await get_account_(name, id, type, page)
    return {'data': data, 'code': 200}


# 删除账号
@app.post('/api/delete_account')
@verify_request_sign
@verify_request_token
@verify_request_permissions
async def delete_account(request: Request, member_id: int = Body(...), id: int = Body(...),
                         modified_by: str = Body(...)):
    data = await delete_account_(id, modified_by)
    return {'data': data, 'code': 200}


# 复制账号
@app.post('/api/copy_account')
@verify_request_sign
@verify_request_token
@verify_request_permissions
async def copy_account(request: Request, member_id: int = Body(...), id: int = Body(...), modified_by: str = Body(...)):
    data = await copy_account_(id, modified_by)
    return {'data': data, 'code': 200}


# 上线触发执行case
@app.get('/api/trigger_execute_case')
async def trigger_execute_case(request: Request, background_tasks: BackgroundTasks, authorization: str, project_id: int,
                               threads: int = 1):
    if authorization != TriggerExecuteToken:
        raise HTTPException(status_code=412, detail="token验证失败")
    # 添加到后台任务
    background_tasks.add_task(trigger_execute_case_, project_id, threads)
    return {'code': 200, 'data': '已开始运行，请稍后通过测试报告页面查看结果'}


if __name__ == '__main__':
    uvicorn.run(
        "app:app",
        host="127.0.0.1",
        port=5004,
        workers=4,
        reload=False,
    )
